import matplotlib.pyplot as plt  
import matplotlib.dates as mdates  
from datetime import datetime  
import argparse

def to_datatime(time_str):
    date_time_str, microsecond_str = time_str.split('.')  
  
    # 确保只取微秒字符串的前六位  
    microsecond_str = microsecond_str[:6]  
    
    # 将微秒字符串转换为整数  
    microseconds = int(microsecond_str)  
    
    # 使用strptime来解析日期和时间部分  
    date_time_obj = datetime.strptime(date_time_str, "%Y-%m-%d %H:%M:%S")  
    
    # 创建一个新的datetime对象（包括微秒）  
    result_datetime = date_time_obj.replace(microsecond=microseconds) 
    return result_datetime

def plot(input_file, type, begin, end):
    data = []
    end = end if end else float('inf')
    for idx, line in enumerate(open(input_file, 'r'), start=1):
        if idx >= begin and idx <= end:
            timestamp, bytes = line.strip().split(': ')  
            tx_bytes, rx_bytes = bytes.strip().split(' ')  
            timestamp = to_datatime(timestamp)
            tx_bytes = int(tx_bytes)  
            rx_bytes = int(rx_bytes)  
            data.append((timestamp, tx_bytes, rx_bytes))  

    rates = []  
    prev_timestamp = None  
    prev_bytes = None  
    
    for timestamp, tx_bytes, rx_bytes in data:  
        if type == 'tx':
            bytes = int(tx_bytes)
        else:
            bytes = int(rx_bytes)
        if prev_timestamp is not None:  
            time_diff = (timestamp - prev_timestamp).total_seconds() # to seconds
            rate = (bytes - prev_bytes) * 8 / time_diff / 10**9 # to Gbps
            
            rates.append((timestamp, rate))  
        prev_timestamp = timestamp  
        prev_bytes = bytes
    
    print('Peak rate: ', max([e[1] for e in rates]))

    fig, ax = plt.subplots()
    if begin is None:
        begin = 1

    ax.plot([e[0] for e in rates], [e[1] for e in rates], label='%s rate (Gbps)' % type)
    # ax.plot([item[0] for item in data[1:]], rx_rates, label='Receive Rate (Gbps)')  
    
    ax.xaxis.set_major_formatter(mdates.DateFormatter('%Y-%m-%d %H:%M:%S'))  
    ax.xaxis.set_major_locator(mdates.MinuteLocator(interval=1))  # 每分钟一个刻度  
    
    plt.xlabel('Time')  
    plt.ylabel('Rate (Gbps)')  
    plt.title('Network Traffic Rates')  
    plt.legend()  
    plt.grid(True)  
    plt.tight_layout()
    # plt.savefig('out.png')
    plt.show()

# plot('2tp.log', 'tx', 11000, 11500)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()  
  
    parser.add_argument('-b', '--begin', default=0, type=int, help='the start line of log')  
    parser.add_argument('-e', '--end', type=int, help='The end line of log')  
    parser.add_argument('-t', '--type', default='tx', help='The traffic direction, tx or rx')
    parser.add_argument('input', type=str, help='The input log file')
    
    args = parser.parse_args() 
    plot(args.input, args.type, args.begin, args.end)